#!usr/bin/env python  
# -*- coding:utf-8 _*-
""" 
@author:robot
@file: tap_connector03_mongodb.py
@version:
@time: 2022/01/13 
"""
import pymongo
from TapClientConnector.connector.tap_connector00_base import *


class TAPMongoDBConnector(TAPConnectorBase):

    def __init__(self, kwargs):
        super().__init__(kwargs)
        user = self.FORMAT_DB_INFO["user"]
        password = self.FORMAT_DB_INFO["password"]
        host = self.FORMAT_DB_INFO["host"].split(':')[0]
        db = self.FORMAT_DB_INFO["database"]
        port = self.FORMAT_DB_INFO["port"]
        if user:
            mongodb_uri = f"mongodb://{user}:{password}@{host}:{port}/{db}?authSource=admin"
        else:
            mongodb_uri = f"mongodb://{host}:{port}/{db}?authSource=admin"
        self.connector = pymongo.MongoClient(mongodb_uri)

    def __mg_drop_db(self, db) -> None:
        self.connector.drop_database(db)

    def __mg_drop_table(self, table) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]].drop_collection(table)

    def __mg_delete_one(self, table, query) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].delete_one(query)

    def __mg_delete_many(self, table, query) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].delete_many(query)

    def __mg_insert_one(self, table, document: dict) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].insert_one(document)

    def __mg_insert_many(self, table, document: list) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].insert_many(document)

    def __mg_update_one(self, table, query, document) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].update_one(query, {'$set': document})

    def __mg_update_many(self, table, query, document) -> None:
        self.connector[self.FORMAT_DB_INFO["database"]][table].update_many(query, {'$set': document})

    def __mg_find(self, table, document=None, show=None) -> list:
        if document is None:
            document = {}
        res = self.connector[self.FORMAT_DB_INFO["database"]][table].find(document, show)
        res_list = list(res)
        return res_list

    @connector_exec_logger
    def _exec(self, sql):
        if isinstance(sql, tuple):
            __op = sql[0]
            __data = sql[1]
            if __op == "insert_one":
                self.__mg_insert_one(*__data)
            elif __op == "insert_many":
                if __data[1] is None:
                    logger.info('mongodb not data exec')
                    return
                self.__mg_insert_many(*__data)
            elif __op == "update_one":
                self.__mg_update_one(*__data)
            elif __op == "update_many":
                self.__mg_update_many(*__data)
            elif __op == "delete_one":
                self.__mg_delete_one(*__data)
            elif __op == "delete_many":
                self.__mg_delete_many(*__data)
            elif __op == "drop_db":
                self.__mg_drop_db(*__data)
            elif __op == "drop_table":
                self.__mg_drop_table(*__data)
            elif __op == "find":
                self.__mg_find(*__data)
            self.connector.close()
        else:
            raise CustomMongoSyntaxError

    def _query_get_column(self, sql) -> list:
        __op = sql[0]
        __data = sql[1]
        if __op == "find":
            response = self.__mg_find(*__data)
            tmp = []
            if response:
                for k, v in response[0].items():
                    tmp.append(k)
            return tmp

    def query_get_column(self, sql):
        result = self._query_get_column(sql)
        self.connector.close()
        return result

    @connector_query_logger
    def query(self, sql):
        result = self._query(sql)
        self.connector.close()
        return result

    def _query(self, sql: tuple):
        """
        (operate, 必填)
        :param sql:
        :return:
        """
        if isinstance(sql, tuple):
            __op = sql[0]
            __data = sql[1]
            if __op == "find":
                return self.__mg_find(*__data)
            elif __op == "show_dbs":
                return self.__show_dbs()
            elif __op == "show_tables":
                return self.__show_tables()
            elif __op == "count":
                return self.__count(*__data)
            else:
                logger.fatal("find | show_dbs | show_tables | count , support")
        else:
            raise CustomMongoSyntaxError

    # **** for mongodb ****

    def __show_dbs(self):
        res = self.connector.list_database_names()
        self.connector.close()
        return res

    def __show_tables(self):
        res = self.connector[self.FORMAT_DB_INFO["database"]].list_collection_names()
        self.connector.close()
        return res

    def __count(self, table, filter: dict) -> list:
        res = self.connector[self.FORMAT_DB_INFO["database"]][table].count_documents(filter)
        self.connector.close()
        return [(res,)]

    @connector_query_logger
    def show_dbs(self):
        return self.__show_dbs()

    @connector_query_logger
    def show_tables(self):
        return self.__show_tables()

    # ***** Generate *****
    # ** DDL **

    @generate_logger
    def generate_create(self, table_column=None, table=None) -> tuple:
        if table_column is None:
            table_column = {'test': 'test'}
        if table:
            self.CURRENT_TABLE = table
        return "insert_many", (self.CURRENT_TABLE, table_column)

    @generate_logger
    def generate_drop(self, table=None):
        if table:
            self.CURRENT_TABLE = table
        return 'drop_table', (self.CURRENT_TABLE,)

    # **DML**

    @generate_logger
    def generate_insert(self, insert_list: list, table=None) -> tuple:
        if table:
            self.CURRENT_TABLE = table
        return "insert_many", (self.CURRENT_TABLE, insert_list)

    @generate_logger
    def generate_update(self, update, where, table=None) -> tuple:
        """
        :param update: {}
        :param where:  {}
        :param table:
        :return:
        """
        if table:
            self.CURRENT_TABLE = table
        return "update_many", (self.CURRENT_TABLE, where, update)

    @generate_logger
    def generate_delete(self, where, table=None):
        if table:
            self.CURRENT_TABLE = table
        if where is None:
            where = {}
        return 'delete_many', (self.CURRENT_TABLE, where)

    @generate_logger
    def generate_select(self, table=None, column='*', where=None, order=''):
        if where is None:
            where = {}
        if table:
            self.CURRENT_TABLE = table
        if column == '*':
            column = {"_id": 0}
        if column == 'count(*)':
            return "count", (self.CURRENT_TABLE, where)
        return "find", (self.CURRENT_TABLE, where, column)


if __name__ == '__main__':
    T = TAPMongoDBConnector(MongoDB.auto_mongo_source)
    res = T.generate_select('mg2')
    print(T.query(res))
    print(T.query_get_column(res))
